/******************************************************************************
 * Project: Processor Affinity Manager                                        *
 * File:    ConfigurationFile.cs                                              *
 * Author:  Aaron Jones                                                       *
 * Version: 1                                                                 *
 *                                                                            *
 * Description:                                                               *
 *    This file contains the implementation for the Windows Service used to   *
 *    apply configured processor affinities.                                  *
 ******************************************************************************/

using System;
using System.Collections.Generic;
using System.Management;
using System.ComponentModel;
using System.Diagnostics;
using System.ServiceProcess;

namespace ProcessorAffinityManager
{
    public partial class Service : ServiceBase
    {
        public const string Name = "Processor Affinity Service";
        public const int IgnoreNextProcessCommand = 128;

        private ManagementEventWatcher m_ProcessEventWatcher;
        private ConfigurationFile m_Config;
        private bool m_IgnoreNextProcess = false;

        public static bool Exists
        {
            get
            {
                try
                {
                    foreach (ServiceController sc in ServiceController.GetServices())
                        if (sc.ServiceName == Name)
                            return true;
                }

                catch
                {
                }

                return false;
            }
        }

        public static bool Running
        {
            get
            {
                try
                {
                    ServiceController sc = new ServiceController(Name);
                    if (sc.Status == ServiceControllerStatus.Running)
                        return true;
                }

                catch
                {
                }

                return false;
            }
        }

        public static void Enable()
        {
            // Set the startmode to automatic.
            string servicePath = "Win32_Service.Name='" + Name + "'";
            ManagementPath mPath = new ManagementPath(servicePath);
            ManagementObject mObj = new ManagementObject(mPath);
            object[] parameters = new object[1];
            parameters[0] = "Automatic";
            mObj.InvokeMethod("ChangeStartMode", parameters); 

            // Start the service
            ServiceController sc = new ServiceController(Name);
            if (sc.Status != ServiceControllerStatus.Running)
                sc.Start();
        }

        public static void Disable()
        {
            // Set the startmode to disabled.
            string servicePath = "Win32_Service.Name='" + Name + "'";
            ManagementPath mPath = new ManagementPath(servicePath);
            ManagementObject mObj = new ManagementObject(mPath);
            object[] parameters = new object[1];
            parameters[0] = "Disabled";
            mObj.InvokeMethod("ChangeStartMode", parameters); 

            // Stop the service.
            ServiceController sc = new ServiceController(Name);
            if (sc.Status == ServiceControllerStatus.Running)
                sc.Stop();
        }

        public Service()
        {
            InitializeComponent();

            m_Config = new ConfigurationFile();

            WqlEventQuery ProcessCreatedQuery = new WqlEventQuery(
                "SELECT * FROM __InstanceCreationEvent WITHIN 2 WHERE TargetInstance ISA 'Win32_Process'");
            m_ProcessEventWatcher = new ManagementEventWatcher(new ManagementScope(), ProcessCreatedQuery);
            m_ProcessEventWatcher.EventArrived += new EventArrivedEventHandler(ProcessCreatedEvent);
        }

        protected override void OnStart(string[] args)
        {
            m_ProcessEventWatcher.Start();
        }

        protected override void OnStop()
        {
            m_ProcessEventWatcher.Stop();
        }

        private void ProcessCreatedEvent(object sender, EventArrivedEventArgs e)
        {
            EventLog.WriteEntry("NEW PROCESS");

            if (m_IgnoreNextProcess)
            {
                m_IgnoreNextProcess = false;
                return;
            }

            try
            {
                ManagementBaseObject mbo = (ManagementBaseObject)e.NewEvent["TargetInstance"];
                string fullPath = (string)mbo["ExecutablePath"];
                int processId = (int)((UInt32)mbo["ProcessId"]);

                EventLog.WriteEntry("NEW PROCESS " + "\"" + fullPath + "\", " + processId.ToString());

                foreach (Default d in m_Config.Defaults)
                {
                    EventLog.WriteEntry("DEFAULT " + d.ExecutablePath + " " + d.ProcessorAffinity.ToString());
                    if (fullPath == d.ExecutablePath)
                    {
                        Process.GetProcessById(processId).ProcessorAffinity = (IntPtr)d.ProcessorAffinity;
                        break;
                    }
                }
            }

            catch (Exception)
            {
                EventLog.WriteEntry("EXCEPTION!");
            }
        }

        protected override void OnCustomCommand(int command)
        {
            if (command == IgnoreNextProcessCommand)
                m_IgnoreNextProcess = true;
        }
    }

    static class ServiceEntry
    {
        /// <summary>
        /// The main entry point for the application.
        /// </summary>
        static void Main()
        {
            ServiceBase[] ServicesToRun;
            ServicesToRun = new ServiceBase[] 
			{ 
				new Service() 
			};
            ServiceBase.Run(ServicesToRun);
        }
    }
}
